"""
用于分段逻辑模型实现
"""
import os
import time

import numpy as np
from osgeo import gdal, osr
from scipy.interpolate import interp1d
from scipy.optimize import curve_fit
from tqdm import tqdm


# 实现logitic模型
# 输入输出示例
def logistic_model(d):
    def logitic(t, a, b, c):
        y = c / (1 + np.exp(a + b * t)) + d
        return y

    return logitic


# 读取tif文件，输出图像数组和仿射变换矩阵
def read_tif(file_path):
    dataset = gdal.Open(file_path)
    # 读成数组
    data = dataset.ReadAsArray()
    # 得到仿射变换矩阵
    geotrans = dataset.GetGeoTransform()
    return data, geotrans


# 写入tif
def write_tif(file_path, data, geotransform, nodata, dataType):
    gdal_type = ''
    if dataType == 'int16':
        gdal_type = gdal.GDT_Int16
    elif dataType == 'int32':
        gdal_type = gdal.GDT_Int32
    elif dataType == 'float32':
        gdal_type = gdal.GDT_Float32
    elif dataType == 'float64':
        gdal_type = gdal.GDT_Float64
    elif dataType == 'byte':
        gdal_type = gdal.GDT_Byte
    elif dataType == 'uint16':
        gdal_type = gdal.GDT_UInt16
    elif dataType == 'uint32':
        gdal_type = gdal.GDT_UInt32
    # 相当于一个创建数据集的驱动
    driver = gdal.GetDriverByName("GTiff")
    # 根据数据的维度来确定行列数（图像大小）
    rows, cols = data.shape
    # 创建一个新的数据集，存储输出文件；1是波段
    dataset = driver.Create(file_path, cols, rows, 1, gdal_type)
    # 设置仿射变换矩阵
    dataset.SetGeoTransform(geotransform)
    # 设置投影
    prj = osr.SpatialReference()
    # WGS84投影
    prj.ImportFromEPSG(4326)
    dataset.SetProjection(prj.ExportToWkt())
    # 1个波段
    band = dataset.GetRasterBand(1)
    band.WriteArray(data)
    band.SetNoDataValue(nodata)
    # 释放内存
    del dataset


# 逐像元插值主函数
def main(input_image_dir, qc_image_dir, pr_image_dir, output_image_dir):
    start = time.time()
    # 先读取所有tif文件
    basename = ".tif"
    evi_list = [os.path.join(input_image_dir, f) for f in os.listdir(input_image_dir) if f.endswith(basename)]
    qc_list = [os.path.join(qc_image_dir, f) for f in os.listdir(qc_image_dir) if f.endswith(basename)]
    pr_list = [os.path.join(pr_image_dir, f) for f in os.listdir(pr_image_dir) if f.endswith(basename)]
    date_list = [os.path.splitext(os.path.basename(f))[0] for f in evi_list]
    print(date_list)
    start_time = int(date_list[0][-3:])
    end_time = int(date_list[-1][-3:])
    print(start_time, end_time)
    time_array = np.arange(start_time, end_time + 1, 16)


    dataset = gdal.Open(evi_list[0])
    rows, cols = dataset.RasterYSize, dataset.RasterXSize
    geotransform = dataset.GetGeoTransform()
    nodata = dataset.GetRasterBand(1).GetNoDataValue()
    # # 存储所有图像为一个列表
    # image_arrays = []
    # for evi_file in evi_list:
    #     data, _ = read_tif(evi_file)
    #     image_arrays.append(data)

    # 读取所有图像
    evi_array = np.zeros((len(evi_list), rows, cols), dtype=np.int16)
    pr_array = np.zeros((len(pr_list), rows, cols), dtype=np.int8)
    qc_array = np.zeros((len(qc_list), rows, cols), dtype=np.uint16)
    for i in range(len(evi_list)):
        data, _ = read_tif(evi_list[i])
        evi_array[i, :, :] = data
    for i in range(len(pr_list)):
        data, _ = read_tif(pr_list[i])
        pr_array[i, :, :] = data
    for i in range(len(qc_list)):
        data, _ = read_tif(qc_list[i])
        # 将质量控制数据中数据转为2进制
        binary_arr = np.vectorize(np.binary_repr)(data)
        # 将上面二维数组每个元素字符串后6位转为十进制组成新的数组
        # decimal_arr = np.array([int(binary[-4:], 2) for binary in binary_arr.flatten()]).reshape(binary_arr.shape)
        # qc_array[i, :, :] = decimal_arr.astype(np.uint8)
        snow_mask = np.char.count(binary_arr, "1", start=1, end=2) == 1
        water_mask = np.char.count(binary_arr, "001", start=2, end=5) != 1
        qc_array[i, :, :] = snow_mask + water_mask

    # 对已有图像进行质量控制,覆盖雪和水体
    # evi_array[evi_array == nodata] = np.NAN
    # evi_array[qc_array == 1] = nodata
    # print(evi_array)

    # 需要插值的日期
    inter_date = []
    for date in date_list:
        while str(int(date) + 1) not in date_list and int(date) + 1 <= int(date_list[-1]):
            inter_date.append(str(int(date) + 1))
            date = str(int(date) + 1)
    print(inter_date)
    inter_time = [int(date[-3:]) for date in inter_date]
    # 创建数组存放插值结果
    result_array = np.zeros((rows, cols, len(inter_date)), dtype=np.int16)

    evi_array = np.transpose(evi_array, (1, 2, 0))
    pr_array = np.transpose(pr_array, (1, 2, 0))
    qc_array = np.transpose(qc_array, (1, 2, 0))

    # 逐像元循环
    for i in tqdm(range(evi_array.shape[0]), desc="逻辑回归"):
        for j in range(evi_array.shape[1]):
            # time_array = np.asarray(date_list).astype(np.int32)
            time_series = evi_array[i, j, :]
            # 获取等于nodata的位置
            # eq_nodata_index = np.where(time_series == nodata)
            # # 删除等于nodata的位置
            # time_series = np.delete(time_series, eq_nodata_index)
            # time_array = np.delete(time_array, eq_nodata_index)
            # if len(time_series) == 0:
            #     # 得到值的同时要得到他的索引，就需要用到枚举
            #     for k, date in enumerate(inter_date):
            #         result_array[i, j, k] = nodata
            #     continue

            # min_evi = np.min(time_series)
            # # 获取logistic模型，并赋值初始参数d
            # logistic = logistic_model(min_evi)
            #
            # # popt拟合所有参数结果的数组，pcov是拟合结果的方差
            # # 最小二乘的曲线拟合
            # popt, pcov = curve_fit(logistic, time_array, time_series, maxfev=100000)
            # a, b, c = popt[0], popt[1], popt[2]
            #
            # # 得到值的同时要得到他的索引，就需要用到枚举
            # for k, date in enumerate(inter_date):
            #     result_array[i, j, k] = logistic(int(date[-3:]), a, b, c)

            # 进行内插
            interp_func = interp1d(time_array, time_series, kind='quadratic')
            result_array[i, j, :] = interp_func(inter_time)

            # print(time_series)
            # print(result_array[i, j, :])
            # print("=" * 20)

    result_array = np.transpose(result_array, (2, 0, 1))
    result_array.dtype = np.int16
    for i in range(result_array.shape[0]):
        output_image = output_image_dir + "/" + inter_date[i] + ".tif"
        write_tif(output_image, result_array[i, :, :], geotransform, nodata=nodata, dataType='int16')  # result_array[i, :, :]中的：代表所有

    #
    # # 逐像元循环
    # for y in range(cols):
    #     for x in range(rows):
    #         # 设置一个列表存放该像元时间序列
    #         time_series_pixels = []
    #         # 循环每个图像存入序列中
    #         for image_array in image_arrays:
    #             time_series_pixels.append(image_array[x, y])
    #
    #         time_series_pixels = np.asarray(time_series_pixels).astype(np.float32)
    #         time_array = np.asarray(date_list).astype(np.int32)
    #         min_evi = np.min(time_series_pixels)
    #         # 获取logistic模型，并赋值初始参数d
    #         logistic = logistic_model(min_evi)
    #
    #         # popt拟合所有参数结果的数组，pcov是拟合结果的方差
    #         # 最小二乘的曲线拟合
    #         popt, pcov = curve_fit(logistic, time_array, time_series_pixels, maxfev=100000)
    #         a, b, c = popt[0], popt[1], popt[2]
    #
    #         # 得到值的同时要得到他的索引，就需要用到枚举
    #         for i, date in enumerate(inter_date):
    #             result_array[i, x, y] = logistic(int(date), a, b, c)
    #
    # for i in range(result_array.shape[0]):
    #     output_image = output_image_dir + "/" + inter_date[i] + ".tif"
    #     write_tif(output_image, result_array[i, :, :], geotransform, nodata=-9999)  # result_array[i, :, :]中的：代表所有
    #
    end = time.time()
    print("完成插值，耗时：{}s".format(end - start))


if __name__ == '__main__':
    # 插值从2013207-2013236
    input_image_dir = r"G:\test\process_result\EVI_tif\merge_clip\2013\EVI"
    qc_image_dir = r"G:\test\process_result\EVI_tif\merge_clip\2013\QC"
    pr_image_dir = r"G:\test\process_result\EVI_tif\merge_clip\2013\PR"
    output_image_dir = r"G:\test\process_result\EVI_tif\EVI_quadratic_2013"
    if not os.path.exists(output_image_dir):
        os.makedirs(output_image_dir)
    main(input_image_dir, qc_image_dir, pr_image_dir, output_image_dir)
